import subprocess
from pycparser import c_parser, c_ast
import json


# --- 1. 预处理 C 代码 ---
def preprocess_code(filename):
    result = subprocess.run(['gcc', '-E', filename], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True,
                            text=True)
    return result.stdout


# --- 2. 将 AST 节点还原为 C 代码字符串 ---
def reconstruct_code(node):
    if isinstance(node, c_ast.Assignment):
        return f"{reconstruct_code(node.lvalue)} {node.op} {reconstruct_code(node.rvalue)}"
    elif isinstance(node, c_ast.ID):
        return node.name
    elif isinstance(node, c_ast.Constant):
        return node.value
    elif isinstance(node, c_ast.BinaryOp):
        return f"({reconstruct_code(node.left)} {node.op} {reconstruct_code(node.right)})"
    elif isinstance(node, c_ast.UnaryOp):
        return f"{node.op}{reconstruct_code(node.expr)}"
    elif isinstance(node, c_ast.FuncCall):
        name = reconstruct_code(node.name)
        args = ', '.join([reconstruct_code(a) for a in node.args.exprs]) if node.args else ''
        return f"{name}({args})"
    elif isinstance(node, c_ast.Cast):
        return f"({reconstruct_code(node.to_type)}){reconstruct_code(node.expr)}"
    elif isinstance(node, c_ast.TypeDecl):
        return node.declname
    elif isinstance(node, c_ast.Decl):
        return node.name
    elif isinstance(node, c_ast.ArrayRef):
        return f"{reconstruct_code(node.name)}[{reconstruct_code(node.subscript)}]"
    elif isinstance(node, c_ast.Compound):
        return '{...}'
    elif isinstance(node, c_ast.Return):
        return f"return {reconstruct_code(node.expr)}"
    elif isinstance(node, c_ast.If):
        return f"if ({reconstruct_code(node.cond)})"
    elif isinstance(node, c_ast.While):
        return f"while ({reconstruct_code(node.cond)})"
    return type(node).__name__


# --- 3. 控制流图与 DU 路径构建类 ---
class CFGBuilder(c_ast.NodeVisitor):
    def __init__(self):
        self.nodes = []  # Store nodes in a list
        self.edges = []  # Store edges in a list
        self.count = 1  # Start node ID from 1
        self.prev = []  # Previous nodes to create edges
        self.definitions = {}  # Track definitions of variables
        self.uses = {}  # Track uses of variables

    def add_node(self, label):
        # 为每个节点添加唯一编号
        node_id = self.count
        self.nodes.append({"id": node_id, "label": label})
        self.count += 1

        # 为当前节点添加边
        if self.prev:
            for p in self.prev:
                self.edges.append({"source": p, "target": node_id})

        self.prev = [node_id]
        return node_id

    def record_definition(self, var_name, node_id):
        if var_name not in self.definitions:
            self.definitions[var_name] = []
        self.definitions[var_name].append(node_id)

    def record_use(self, var_name, node_id):
        if var_name not in self.uses:
            self.uses[var_name] = []
        self.uses[var_name].append(node_id)

    def visit_Assignment(self, node):
        # Record the definition of the variable
        var_name = reconstruct_code(node.lvalue)  # Get the left-hand side variable
        node_id = self.add_node(f"Assignment: {reconstruct_code(node)}")
        self.record_definition(var_name, node_id)

    def visit_ID(self, node):
        # Record the use of the variable
        var_name = node.name
        node_id = self.prev[-1]  # Use the last node as the position of usage
        self.record_use(var_name, node_id)

    def visit_Decl(self, node):
        if node.init:
            # If there's an initializer, it's a definition
            var_name = node.name
            node_id = self.add_node(f"Decl: {reconstruct_code(node)}")
            self.record_definition(var_name, node_id)

    def visit_FuncCall(self, node):
        # For function calls, we can track arguments as uses
        node_id = self.add_node(f"FuncCall: {reconstruct_code(node)}")
        if node.args:
            for arg in node.args.exprs:
                self.visit(arg)  # Visit arguments to track their uses

    def visit_Return(self, node):
        node_id = self.add_node(f"Return: {reconstruct_code(node)}")
        if isinstance(node.expr, c_ast.ID):
            self.record_use(node.expr.name, node_id)

    def visit_If(self, node):
        cond = reconstruct_code(node.cond)
        cond_id = self.add_node(f"If: {cond}")
        self.record_use(cond, cond_id)

        prev_copy = self.prev[:]
        self.prev = [cond_id]
        self.visit(node.iftrue)
        true_exit = self.prev[:]

        self.prev = [cond_id]
        if node.iffalse:
            self.visit(node.iffalse)
            false_exit = self.prev[:]
        else:
            false_exit = [cond_id]

        self.prev = true_exit + false_exit

    def visit_While(self, node):
        cond = reconstruct_code(node.cond)
        cond_id = self.add_node(f"While: {cond}")
        self.record_use(cond, cond_id)

        loop_start = cond_id
        self.prev = [cond_id]
        self.visit(node.stmt)
        if self.prev:
            for p in self.prev:
                self.edges.append({"source": p, "target": loop_start})
        self.prev = [cond_id]

    def generate_du_paths(self):
        du_paths = []
        for var_name, defs in self.definitions.items():
            if var_name in self.uses:
                for def_node in defs:
                    for use_node in self.uses[var_name]:
                        du_paths.append({
                            "variable": var_name,
                            "definition_node": def_node,
                            "use_node": use_node
                        })
        return du_paths


# --- 4. 生成控制流图和 DU 路径的 JSON ---
def generate_cfg_json_with_du(c_code):
    parser = c_parser.CParser()
    try:
        ast = parser.parse(c_code)
    except Exception as e:
        return {"status": 515, "message": "Error", "data": str(e)}

    builder = CFGBuilder()
    builder.visit(ast)

    du_paths = builder.generate_du_paths()

    demo_json = {
        "nodes": builder.nodes,
        "edges": builder.edges,
        "du_paths": du_paths
    }
    return {"status": 200, "message": "Success", "data": json.dumps(demo_json, indent=4)}


# --- 测试代码 ---
c_code = """
int main() {
    int a = 5;
    if (a > 3) {
        a = a - 1;
    } else {
        a = a + 1;
    }
    return a;
}
"""

cfg_json = generate_cfg_json_with_du(c_code)
print(cfg_json.get)
